from crf import CRFModel
from evaluating import Metrics
import os, sys, pickle

CRF_MODEL_PATH = './model/crf.pkl'


def build_corpus(model, make_vocab=True, data_dir="./data"):
    """读取数据"""
    assert model in ['train', 'test']

    word_lists = []
    tag_lists = []
    with open(os.path.join(data_dir, model + ".char.txt"), 'r', encoding='utf-8') as f:
        word_list = []
        tag_list = []
        for line in f:
            if line != '\n':
                word, tag = line.strip('\n').split()
                word_list.append(word)
                tag_list.append(tag)
            else:
                word_lists.append(word_list)
                tag_lists.append(tag_list)
                word_list = []
                tag_list = []

    # 如果make_vocab为True，还需要返回word2id和tag2id
    if make_vocab:
        word2id = build_map(word_lists)
        tag2id = build_map(tag_lists)

        return word_lists, tag_lists, word2id, tag2id
    else:
        # word2idt = build_map(word_lists)
        # tag2idt = build_map(tag_lists)
        # print(word2idt)
        # print(tag2idt)
        '''for e in range(len(tag_lists)):
            for w in range(len(tag_lists[e])):
                print(tag_lists[e][w])
                print("\n")
        q = []
        z = 0
        for index, value in enumerate(tag_lists):
            for q, w in enumerate(tag_lists[index]):

                    for m, n in enumerate(word_lists):
                       for o, l in enumerate(word_lists[m]):
                           if w != 'O':
                    #print(q)
                                k = 0
                                if o == q and k == 0:
                                    print(o)
                                    k = 1'''

        return word_lists, tag_lists


def build_map(lists):
    maps = {}
    for list_ in lists:
        for e in list_:
            if e not in maps:
                maps[e] = len(maps)
    # print(maps)
    return maps


def crf_train_eval(train_data, test_data, remove_O=False):
    # 训练CRF模型
    train_word_lists, train_tag_lists = train_data
    test_word_lists, test_tag_lists = test_data

    crf_model = CRFModel()
    crf_model.train(train_word_lists, train_tag_lists)
    save_model(crf_model, CRF_MODEL_PATH)

    pred_tag_lists = crf_model.test(test_word_lists)

    metrics = Metrics(test_tag_lists, pred_tag_lists, remove_O=remove_O)
    metrics.report_scores()
    '''for b in range(len(test_tag_lists)):
        print(test_tag_lists[b])
        print("\n")'''
    TP = 0
    FN = 0
    TN = 0
    FP = 0
    for c in range(len(test_tag_lists)):
        for d in range(len(test_tag_lists[c])):
            if test_tag_lists[c][d] == pred_tag_lists[c][d]:
                TP = TP + 1
            else:
                FN = FN + 1
    Precision = TP / (TP + FP)
    Accuracy = TP / (TP + FN + TN + FP)
    Recall = TP / (TP + FN)
    f = 2 * TP / (2 * TP + FP + FN)
    print("\n")
    print("精确率为：", Precision)
    print("准确率：", Accuracy)
    print("召回率：", Recall)
    print("F值：", f)
    return pred_tag_lists


def save_model(model, file_name):
    """用于保存模型"""
    with open(file_name, "wb") as f:
        pickle.dump(model, f)


def main():
    """训练模型，评估结果"""

    # 读取数据
    with open('data/test.char.txt', 'r', encoding='utf-8') as f:
        # line = f.readline().strip()
        l = 0
        for line in f:
            l = l
            if line != '\n':
                word1, tag1 = line.strip('\n').split()
                if tag1 != 'O':
                    l = 0
                    print(word1)
                elif tag1 == 'O' and l == 0:
                    l = 1
                    print(' ')
    print('loading...')
    train_word_lists, train_tag_lists, word2id, tag2id = build_corpus("train")
    # test_word_lists, test_tag_lists = build_corpus("test", make_vocab=False)
    test_word_lists, test_tag_lists = build_corpus("test", make_vocab=False)
    # 训练评估CRF模型
    print('training...')
    crf_pred = crf_train_eval((train_word_lists, train_tag_lists), (test_word_lists, test_tag_lists))


if __name__ == "__main__":
    main()
